import numpy as np
import scipy
from scipy.fft import dst,idst


def DST(a,*args):
    shape = a.shape
    if len(shape) == 1:
        if shape[0] > 1:
            do_trans = 1
        else:
            do_trans = 0
        a = np.ravel(a)
        a = a.reshape(-1,1)
    else:
        do_trans = 0

    shape = a.shape

    if args:
        n = args[0]
    else:
        n = shape[0]

    m = a.shape[1]

    if shape[0] < n:
        aa = np.zeros((n,m))
        aa[0:shape[0],:] = a
    else:
        aa = a[0:n,:]

    y = np.zeros((2 * (n + 1),m))
    y[1:n+1,:] = aa

    y[n+2:2*(n+1),:] = -np.flipud(aa)
    yy = np.fft.fft(y,axis=0)

    b = yy[1:n+1,:] / (-2j)

    if np.isreal(a.any()):
        b = np.real(b)
    if do_trans:
        b = b.T

    return b

def IDST(a,*args):
    if args:
        n = args[0]
    else:
        n = a.shape[0]

    nn = n + 1
    b = 2/nn * DST(a,n)
    return b

def dst2(x):
    x_dst_rows=dst(x,axis=0,norm='ortho')
    x_dst_2d=dst(x_dst_rows,axis=1,norm='ortho')
    #return scipy.fftpack.dst(scipy.fftpack.dst(x.T, norm='ortho').T, norm='ortho')
    return x_dst_2d

# 二维离散余弦反变换
def idst2(x):
    x_idst_rows=idst(x,axis=1,norm='ortho')
    x_idst_2d=idst(x_idst_rows,axis=0,norm='ortho')
    return x_idst_2d
    #return scipy.fftpack.idst(scipy.fftpack.idst(x.T, norm='ortho').T, norm='ortho')

a = np.array(range(1,13)).reshape(-1, 3).T

# x = [1,3,5,3,1,-1,-3,-5,-3,-1]
# print(dst(x))


print(DST(a))
#print(dst2(a))